import os
import tqdm

root = "/data/ct/gunman/datasets/COD10k"
for dataset in ['CAMO', "CHAMELEON", "COD10K"]:
    train_folder = os.path.join('TestDataset', dataset)
    output_file = f'/data/ct/gunman/UGTR/train_test_file/{dataset}_test_1.lst'

    with open(output_file, 'w') as f:
        for Imgfilename in tqdm.tqdm(os.listdir(f"{os.path.join(root, train_folder)}/Imgs")):
            Imgfilename = f"{train_folder}/Imgs/{Imgfilename}"
            GTfilename = f"{train_folder}/GT/{Imgfilename.split('/')[-1].split('.')[0]}.png"
            Edgefilename = f"{train_folder}/Edge/{Imgfilename.split('/')[-1].split('.')[0]}.png"
            # 判断GT和Edge是否存在
            if os.path.exists(f"{os.path.join(root, GTfilename)}") and os.path.exists(f"{os.path.join(root, Edgefilename)}"):
                f.write(f"{Imgfilename} {GTfilename} {Edgefilename}\n")